package com.diao.redis;

import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;
import redis.clients.jedis.params.SetParams;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeoutException;
/**
 * @author chenzhidiao
 * @version 1.0
 * @date 2020/5/1 19:06
 * @description:
 */
public class RedisDistributeLockUtil {
    private static final String LOCK_SUCCESS = "OK";
    private JedisPool jedisPool;
    private final Map<String, Boolean> DAEMON_MAP = new ConcurrentHashMap<String, Boolean>();

    /**
     * 以下是Lua脚本，通过jedis的eval()方法将Lua代码交给Redis服务端执行。eval命令执行Lua代码的时候，Lua脚本将被当成一个命令去执行，
     * 并且直到eval命令执行完成，Redis才会执行其他命令，从而保证原子性
     */
    /**
     * 如果当前key有值(即有锁)，且值是当前请求的id，那么更新key的过期时间，返回1
     * 如果当前key有值但不是当前请求的id，返回0
     * 如果当前key无值，用当前请求的id为值创建锁，返回1
     */
    private static final String REENTRANT_LOCK_SCRIPT = "local value = redis.call('get',KEYS[1]) " +
            "if value then " +
            "  if value == ARGV[1] then  " +
            "    return redis.call('pexpire', KEYS[1], ARGV[2])  " +
            "  else  " +
            "    return 0  " +
            "  end " +
            "else " +
            "  redis.call('set',KEYS[1],ARGV[1],'NX','PX',ARGV[2]) " +
            "  return 1 " +
            "end";
    /**
     * 如果当前key有值(即有锁)，且值是当前请求的id，那么更新key的过期时间，返回1
     * 如果当前key有值但不是当前请求的id，返回剩余时间
     * 如果当前key无值，用当前请求的id为值创建锁，返回1
     */
    private static final String WAIT_LOCK_SCRIPT = "local value = redis.call('get',KEYS[1]) " +
            "if value then " +
            "  if value == ARGV[1] then  " +
            "    return redis.call('pexpire', KEYS[1], ARGV[2])  " +
            "  else  " +
            "    return redis.call('pttl', KEYS[1])  " +
            "  end " +
            "else " +
            "  redis.call('set',KEYS[1],ARGV[1],'NX','PX',ARGV[2]) " +
            "  return 1 " +
            "end";
    /**
     * 如果当前key有值(即有锁)，且值是当前请求的id，那么更新key的过期时间，返回1
     * 否则返回0
     */
    private static final String REFRESH_LOCK_SCRIPT = "if redis.call('get', KEYS[1]) == ARGV[1] then " +
            "  return redis.call('pexpire', KEYS[1],ARGV[2])  " +
            "else " +
            "  return 0 " +
            "end";
    /**
     * 如果当前key有值(即有锁)，且值是当前请求的id，删除锁，返回1
     * 否则返回0
     */
    private static final String RELEASE_LOCK_SCRIPT = "if redis.call('get', KEYS[1]) == ARGV[1] then " +
            "  return redis.call('del', KEYS[1]) " +
            "else " +
            "  return 0 " +
            "end";

    public RedisDistributeLockUtil(JedisPool jedisPool) {
        this.jedisPool = jedisPool;
    }

    protected Jedis getJedis() {
        return jedisPool.getResource();
    }

    /**
     * 获取redis分布式锁，只要锁存在就返回false
     *
     * @param lockKey   锁的名字
     * @param lockOwner 锁的所有者,通常为当前请求的requestID
     * @param expire    锁过期时间单位毫秒
     * @return
     */
    public boolean getSimpleRedisLock(String lockKey, String lockOwner, Long expire) {
        SetParams setParams = SetParams.setParams().nx().px(expire);

        String result = getJedis().set(lockKey, lockOwner, setParams);
        return result != null ? LOCK_SUCCESS.equals(result) : false;
    }

    /**
     * 获取redis分布式锁, 如果锁存在且属于当前申请者，则刷新锁过期时间，返回true
     *
     * @param lockKey   锁的名字
     * @param lockOwner 锁的所有者,通常为当前请求的requestID
     * @param expire    锁过期时间单位毫秒
     * @return
     */
    public boolean getReentrantRedisLock(String lockKey, String lockOwner, Long expire) {
        List keyList = Collections.singletonList(lockKey);
        List argvList = Arrays.asList(lockOwner,expire.toString());
        Long result = (Long) getJedis().eval(REENTRANT_LOCK_SCRIPT, keyList, argvList);
        return result > 0;
    }


    /**
     * 获取Redis锁，如果锁被占用则阻塞
     *
     * @param lockKey   锁的名字
     * @param lockOwner 锁的所有者,通常为当前请求的requestID
     * @param expire    锁过期时间单位毫秒
     * @param timeOut   超时时间毫秒
     * @return
     * @throws InterruptedException
     * @throws TimeoutException
     */
    public boolean waitRedisLock(String lockKey, String lockOwner, Long expire, Long timeOut) throws InterruptedException, TimeoutException {
        long startTime = System.currentTimeMillis();
        List keyList = Collections.singletonList(lockKey);
        List argvList = Arrays.asList(lockOwner,expire.toString());
        while (true) {
            long nowTime = System.currentTimeMillis();
            if (timeOut > 0 && (nowTime - startTime) >= timeOut) {
                throw new TimeoutException();//超时
            }
            try {
                Long result = (Long) getJedis().eval(WAIT_LOCK_SCRIPT, keyList, argvList);
                if (result == 1) {
                    return true;
                }
                //未获取锁，休眠该锁对应的过期时间，然后继续尝试获取
                Thread.sleep(result);
            } catch (Exception e) {
                Thread.sleep((timeOut >> 8) + 10);
            }
        }
    }

    /**
     * 刷新锁过期时间
     *
     * @param lockKey
     * @param lockOwner
     * @param expire
     */
    public boolean refreshLockExpire(String lockKey, String lockOwner, Long expire) {
        List keyList = Collections.singletonList(lockKey);
        List argvList = Arrays.asList(lockOwner,expire.toString());
        return (Long) getJedis().eval(REFRESH_LOCK_SCRIPT, keyList, argvList) > 0;
    }

    /**
     * 释放锁
     *
     * @param lockKey
     * @param lockOwner
     * @return
     */
    public boolean releaseLock(String lockKey, String lockOwner) {
        Long result = (Long) getJedis().eval(RELEASE_LOCK_SCRIPT, Collections.singletonList(lockKey), Collections.singletonList(lockOwner));
        if (result > 0) {
            DAEMON_MAP.remove(lockKey);
            return true;
        }
        return false;
    }

    /**
     * 获取Redis锁，如果锁被占用则阻塞,启动一个守护线程用来保证在任务进行中锁不过期
     *
     * @param lockKey   锁的名字
     * @param lockOwner 锁的所有者,通常为当前请求的requestID
     * @param expire    锁过期时间单位毫秒
     * @return
     */
    public boolean waitRedisLockWithDaemon(final String lockKey, final String lockOwner, final Long expire, final long timeOut) throws TimeoutException, InterruptedException {
        boolean sig = waitRedisLock(lockKey, lockOwner, expire, timeOut);
        if (sig) {
            DAEMON_MAP.put(lockKey, true);
            lockDaemon(lockKey, lockOwner, expire);
        }
        return sig;
    }

    /**
     * 获取redis分布式锁,启动一个守护线程用来保证在任务进行中锁不过期
     *
     * @param lockKey   锁的名字
     * @param lockOwner 锁的所有者,通常为当前请求的requestID
     * @param expire    锁过期时间单位毫秒
     * @return
     */
    public boolean getRedisLockWithDaemon(final String lockKey, final String lockOwner, final Long expire) {
        boolean sig = getSimpleRedisLock(lockKey, lockOwner, expire);
        if (sig) {
            DAEMON_MAP.put(lockKey, true);
            lockDaemon(lockKey, lockOwner, expire);
        }
        return sig;
    }

    /**
     * 锁的守护线程，保证在锁持有期内不会过期
     *
     * @param lockKey
     * @param lockOwner
     * @param expire
     */
    protected void lockDaemon(final String lockKey, final String lockOwner, final Long expire) {
        long exp;
        if (expire > 1000) {
            exp = expire >> 1;
        } else {
            exp = expire-5;
        }
        final long finalExp = exp;
        Thread thread = new Thread(new Runnable() {
            public void run() {
                Boolean sig = DAEMON_MAP.get(lockKey);
                if (sig == null) {
                    sig = false;
                }
                try {
                    while (sig) {
                        if (refreshLockExpire(lockKey, lockOwner, expire)) {
                            Thread.sleep(finalExp);
                            sig = DAEMON_MAP.get(lockKey);
                        } else {
                            sig = false;
                            DAEMON_MAP.remove(lockKey);
                        }
                    }
                } catch (Exception e) {
                    throw new RuntimeException(e);
                }
            }
        });
        thread.setDaemon(true);
        thread.start();
    }
}
